/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file batch_layout.h
 * \brief
 */
#ifndef IMPL_MATMUL_MODULES_STAGE_COPY_CUBE_IN_BATCH_LAYOUT_H
#define IMPL_MATMUL_MODULES_STAGE_COPY_CUBE_IN_BATCH_LAYOUT_H

namespace matmul {
template<typename IMPL, class INPUT_TYPE, typename = void>
class BatchLayout
{
    MATMUL_USE_MODULE_ON(MatmulVar, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(CopyCubeInParams, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeTiling, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeInfo, INPUT_TYPE::TAG);
    using SrcT = typename INPUT_TYPE::T;
public:
    __aicore__ inline bool IsLayoutGValid()
    {
        return true;
    }

    __aicore__ inline uint32_t GetLoopNum()
    {
        return 1;
    }

    __aicore__ inline uint32_t GetBatchNum()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum();
    }

    __aicore__ inline uint32_t GetSrcStride(const uint32_t stride, const uint32_t height, const uint32_t width)
    {
        if constexpr (PhyPosIsUB(INPUT_TYPE::pos)) {
            return stride != 0 ? stride : CeilAlign<int32_t>(height, BLOCK_CUBE) * CeilAlign<int32_t>(width, c0Size_);
        } else {
            return stride != 0 ? stride : height * width;
        }
    }

    template <bool IS_TRANS>
    __aicore__ inline int32_t GetSrcDValue()
    {
        return MATMUL_MODULE(MatmulShapeInfo)->template GetSingleWidth<IS_TRANS>();
    }

private:
    constexpr static int32_t c0Size_ = AuxGetC0Size<SrcT>();
};

template<typename IMPL, class INPUT_TYPE>
class BatchLayout<IMPL, INPUT_TYPE, enable_if_t<(INPUT_TYPE::layout == LayoutMode::BSNGD)>>
{
    MATMUL_USE_MODULE_ON(MatmulVar, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeInfo, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(CopyCubeInParams, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeTiling, INPUT_TYPE::TAG);
    using SrcT = typename INPUT_TYPE::T;

    __aicore__ inline int32_t GetMaxLayoutInfoG() const
    {
        return MATMUL_CONST_PARAM_VAR.tiling_.GetALayoutInfoG() > MATMUL_CONST_PARAM_VAR.tiling_.GetBLayoutInfoG() ?
               MATMUL_CONST_PARAM_VAR.tiling_.GetALayoutInfoG() : MATMUL_CONST_PARAM_VAR.tiling_.GetBLayoutInfoG();
    }
public:
    __aicore__ inline bool IsLayoutGValid()
    {
        return (MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum() <=
               (MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoN() * GetMaxLayoutInfoG()));
    }

    __aicore__ inline uint32_t GetLoopNum()
    {
        return Ceil(MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum(), GetLayoutInfoNG());
    }

    __aicore__ inline uint32_t GetBatchNum()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum() < GetLayoutInfoNG() ?
            MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum() : GetLayoutInfoNG();
    }

    __aicore__ inline uint32_t GetSrcStride(const uint32_t stride, const uint32_t height, const uint32_t width)
    {
        if constexpr (PhyPosIsUB(INPUT_TYPE::pos)) {
            return stride != 0 ? stride : CeilAlign<int32_t>(width, c0Size_);
        } else {
            return stride != 0 ? stride : width;
        }
    }

    template <bool IS_TRANS>
    __aicore__ inline int32_t GetSrcDValue()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoD() *
            MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoN() *
            MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoG();
    }
private:
    constexpr static int32_t c0Size_ = AuxGetC0Size<SrcT>();

private:
    __aicore__ inline int32_t GetLayoutInfoNG()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoN() * MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoG();
    }
};

template<typename IMPL, class INPUT_TYPE>
class BatchLayout<IMPL, INPUT_TYPE, enable_if_t<(INPUT_TYPE::layout == LayoutMode::SBNGD)>>
{
    MATMUL_USE_MODULE_ON(MatmulVar, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeInfo, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(CopyCubeInParams, INPUT_TYPE::TAG);
    MATMUL_USE_MODULE_ON(MatmulShapeTiling, INPUT_TYPE::TAG);
    using SrcT = typename INPUT_TYPE::T;
public:
    __aicore__ inline bool IsLayoutGValid()
    {
        return true;
    }

    __aicore__ inline uint32_t GetLoopNum()
    {
        return 1;
    }

    __aicore__ inline uint32_t GetBatchNum()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetBatchNum();
    }

    __aicore__ inline uint32_t GetSrcStride(const uint32_t stride, const uint32_t height, const uint32_t width)
    {
        if constexpr (PhyPosIsUB(INPUT_TYPE::pos)) {
            return stride != 0 ? stride : CeilAlign<int32_t>(width, c0Size_);
        } else {
            return stride != 0 ? stride : width;
        }
    }

    template <bool IS_TRANS>
    __aicore__ inline int32_t GetSrcDValue()
    {
        return MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoD() *
            MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoN() *
            MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoG() *
            MATMUL_MODULE(MatmulShapeTiling)->GetLayoutInfoB();
    }
private:
    constexpr static int32_t c0Size_ = AuxGetC0Size<SrcT>();
};

}      // namespace matmul
#endif // IMPL_MATMUL_MODULES_STAGE_COPY_CUBE_IN_BATCH_LAYOUT_H